package com.zenithmind.chat.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.zenithmind.chat.pojo.dto.ChatMessageDTO;
import com.zenithmind.chat.pojo.dto.ChatPrivateMessageDTO;
import com.zenithmind.chat.pojo.vo.ChatMessageVO;
import com.zenithmind.chat.pojo.vo.ChatPrivateMessageVO;
import com.zenithmind.chat.service.ChatMessageService;
import com.zenithmind.chat.service.ChatPrivateMessageService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * 聊天WebSocket处理器
 * 
 * @author ZenithMind Team
 * @since 2025-01-09
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class ChatWebSocketHandler implements WebSocketHandler {

    private final ChatMessageService messageService;
    private final ChatPrivateMessageService privateMessageService;
    private final ObjectMapper objectMapper;

    // 存储所有WebSocket连接 - 用户ID -> WebSocket会话集合
    private static final Map<String, CopyOnWriteArraySet<WebSocketSession>> USER_SESSIONS = new ConcurrentHashMap<>();
    
    // 存储群组订阅关系 - 群组ID -> 用户ID集合
    private static final Map<String, CopyOnWriteArraySet<String>> GROUP_SUBSCRIBERS = new ConcurrentHashMap<>();
    
    // 存储会话信息 - 会话ID -> 用户信息
    private static final Map<String, UserSessionInfo> SESSION_INFO = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String userId = (String) session.getAttributes().get("userId");
        String groupId = (String) session.getAttributes().get("groupId");
        String clientIp = (String) session.getAttributes().get("clientIp");
        
        // 存储会话信息
        UserSessionInfo sessionInfo = new UserSessionInfo();
        sessionInfo.setUserId(userId);
        sessionInfo.setGroupId(groupId);
        sessionInfo.setClientIp(clientIp);
        sessionInfo.setConnectTime(System.currentTimeMillis());
        SESSION_INFO.put(session.getId(), sessionInfo);
        
        // 添加用户会话
        USER_SESSIONS.computeIfAbsent(userId, k -> new CopyOnWriteArraySet<>()).add(session);
        
        // 订阅群组
        if (StringUtils.hasText(groupId)) {
            GROUP_SUBSCRIBERS.computeIfAbsent(groupId, k -> new CopyOnWriteArraySet<>()).add(userId);
        }
        
        log.info("WebSocket连接建立：用户ID={}, 群组ID={}, 会话ID={}, IP={}", 
                userId, groupId, session.getId(), clientIp);
        
        // 发送连接成功消息
        sendMessage(session, createSystemMessage("CONNECTION_SUCCESS", "连接成功"));
        
        // 通知群组其他成员用户上线
        if (StringUtils.hasText(groupId)) {
            broadcastToGroup(groupId, createSystemMessage("USER_ONLINE", userId + " 上线了"), userId);
        }
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        if (!(message instanceof TextMessage)) {
            return;
        }
        
        String payload = ((TextMessage) message).getPayload();
        log.debug("收到WebSocket消息：{}", payload);
        
        try {
            // 解析消息
            WebSocketRequest wsRequest = objectMapper.readValue(payload, WebSocketRequest.class);
            String messageType = wsRequest.getType();
            
            UserSessionInfo sessionInfo = SESSION_INFO.get(session.getId());
            if (sessionInfo == null) {
                sendMessage(session, createErrorMessage("SESSION_NOT_FOUND", "会话信息不存在"));
                return;
            }
            
            switch (messageType) {
                case "SEND_MESSAGE":
                    handleSendMessage(session, wsRequest, sessionInfo);
                    break;
                case "SEND_PRIVATE_MESSAGE":
                    handleSendPrivateMessage(session, wsRequest, sessionInfo);
                    break;
                case "JOIN_GROUP":
                    handleJoinGroup(session, wsRequest, sessionInfo);
                    break;
                case "LEAVE_GROUP":
                    handleLeaveGroup(session, wsRequest, sessionInfo);
                    break;
                case "TYPING":
                    handleTyping(session, wsRequest, sessionInfo);
                    break;
                case "READ_MESSAGE":
                    handleReadMessage(session, wsRequest, sessionInfo);
                    break;
                case "READ_PRIVATE_MESSAGE":
                    handleReadPrivateMessage(session, wsRequest, sessionInfo);
                    break;
                case "HEARTBEAT":
                    handleHeartbeat(session, sessionInfo);
                    break;
                default:
                    sendMessage(session, createErrorMessage("UNKNOWN_MESSAGE_TYPE", "未知消息类型"));
            }
        } catch (Exception e) {
            log.error("处理WebSocket消息异常", e);
            sendMessage(session, createErrorMessage("MESSAGE_PROCESS_ERROR", "消息处理异常"));
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输异常：会话ID={}", session.getId(), exception);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        UserSessionInfo sessionInfo = SESSION_INFO.remove(session.getId());
        if (sessionInfo == null) {
            return;
        }
        
        String userId = sessionInfo.getUserId();
        String groupId = sessionInfo.getGroupId();
        
        // 移除用户会话
        CopyOnWriteArraySet<WebSocketSession> userSessions = USER_SESSIONS.get(userId);
        if (userSessions != null) {
            userSessions.remove(session);
            if (userSessions.isEmpty()) {
                USER_SESSIONS.remove(userId);
            }
        }
        
        // 取消群组订阅
        if (StringUtils.hasText(groupId)) {
            CopyOnWriteArraySet<String> subscribers = GROUP_SUBSCRIBERS.get(groupId);
            if (subscribers != null) {
                // 检查用户是否还有其他连接在该群组
                boolean hasOtherConnection = userSessions != null && userSessions.stream()
                    .anyMatch(s -> groupId.equals(SESSION_INFO.get(s.getId()).getGroupId()));
                
                if (!hasOtherConnection) {
                    subscribers.remove(userId);
                    if (subscribers.isEmpty()) {
                        GROUP_SUBSCRIBERS.remove(groupId);
                    }
                    
                    // 通知群组其他成员用户下线
                    broadcastToGroup(groupId, createSystemMessage("USER_OFFLINE", userId + " 下线了"), userId);
                }
            }
        }
        
        log.info("WebSocket连接关闭：用户ID={}, 群组ID={}, 会话ID={}, 状态={}", 
                userId, groupId, session.getId(), closeStatus);
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 处理发送消息
     */
    private void handleSendMessage(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        try {
            ChatMessageDTO messageDTO = objectMapper.convertValue(wsRequest.getData(), ChatMessageDTO.class);
            
            // 发送消息
            String messageId = messageService.sendMessage(messageDTO, sessionInfo.getUserId());
            
            // 获取完整消息信息
            ChatMessageVO messageVO = messageService.getMessageById(messageId, sessionInfo.getUserId());
            
            // 广播消息到群组
            broadcastToGroup(messageDTO.getGroupId(), createDataMessage("NEW_MESSAGE", messageVO), null);
            
            // 发送发送成功确认
            sendMessage(session, createDataMessage("MESSAGE_SENT", Map.of("messageId", messageId)));
            
        } catch (Exception e) {
            log.error("处理发送消息异常", e);
            sendMessage(session, createErrorMessage("SEND_MESSAGE_ERROR", "发送消息失败"));
        }
    }

    /**
     * 处理加入群组
     */
    private void handleJoinGroup(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        String groupId = (String) wsRequest.getData();
        if (!StringUtils.hasText(groupId)) {
            sendMessage(session, createErrorMessage("INVALID_GROUP_ID", "群组ID无效"));
            return;
        }
        
        // 更新会话信息
        sessionInfo.setGroupId(groupId);
        
        // 订阅群组
        GROUP_SUBSCRIBERS.computeIfAbsent(groupId, k -> new CopyOnWriteArraySet<>()).add(sessionInfo.getUserId());
        
        // 发送加入成功消息
        sendMessage(session, createSystemMessage("GROUP_JOINED", "已加入群组"));
        
        // 通知群组其他成员
        broadcastToGroup(groupId, createSystemMessage("USER_JOINED", sessionInfo.getUserId() + " 加入了群组"), sessionInfo.getUserId());
    }

    /**
     * 处理离开群组
     */
    private void handleLeaveGroup(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        String groupId = sessionInfo.getGroupId();
        if (!StringUtils.hasText(groupId)) {
            return;
        }
        
        // 取消订阅
        CopyOnWriteArraySet<String> subscribers = GROUP_SUBSCRIBERS.get(groupId);
        if (subscribers != null) {
            subscribers.remove(sessionInfo.getUserId());
            if (subscribers.isEmpty()) {
                GROUP_SUBSCRIBERS.remove(groupId);
            }
        }
        
        // 更新会话信息
        sessionInfo.setGroupId(null);
        
        // 发送离开成功消息
        sendMessage(session, createSystemMessage("GROUP_LEFT", "已离开群组"));
        
        // 通知群组其他成员
        broadcastToGroup(groupId, createSystemMessage("USER_LEFT", sessionInfo.getUserId() + " 离开了群组"), sessionInfo.getUserId());
    }

    /**
     * 处理正在输入
     */
    private void handleTyping(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        String groupId = sessionInfo.getGroupId();
        if (!StringUtils.hasText(groupId)) {
            return;
        }
        
        // 广播正在输入状态
        Map<String, Object> typingData = Map.of(
            "userId", sessionInfo.getUserId(),
            "isTyping", wsRequest.getData()
        );
        broadcastToGroup(groupId, createDataMessage("USER_TYPING", typingData), sessionInfo.getUserId());
    }

    /**
     * 处理消息已读
     */
    private void handleReadMessage(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        try {
            @SuppressWarnings("unchecked")
            Map<String, Object> data = (Map<String, Object>) wsRequest.getData();
            @SuppressWarnings("unchecked")
            java.util.List<String> messageIds = (java.util.List<String>) data.get("messageIds");
            
            // 标记消息为已读
            messageService.markMessagesAsRead(messageIds, sessionInfo.getUserId());
            
            // 发送已读确认
            sendMessage(session, createSystemMessage("MESSAGES_READ", "消息已标记为已读"));
            
        } catch (Exception e) {
            log.error("处理消息已读异常", e);
            sendMessage(session, createErrorMessage("READ_MESSAGE_ERROR", "标记已读失败"));
        }
    }

    /**
     * 处理发送私聊消息
     */
    private void handleSendPrivateMessage(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        try {
            ChatPrivateMessageDTO messageDTO = objectMapper.convertValue(wsRequest.getData(), ChatPrivateMessageDTO.class);

            // 发送私聊消息
            String messageId = privateMessageService.sendPrivateMessage(messageDTO, sessionInfo.getUserId());

            // 获取完整消息信息
            ChatPrivateMessageVO messageVO = privateMessageService.getMessageById(messageId, sessionInfo.getUserId());

            // 发送给接收者
            sendToUser(messageDTO.getReceiverId(), createDataMessage("NEW_PRIVATE_MESSAGE", messageVO));

            // 发送发送成功确认给发送者
            sendMessage(session, createDataMessage("PRIVATE_MESSAGE_SENT", Map.of("messageId", messageId)));

        } catch (Exception e) {
            log.error("处理发送私聊消息异常", e);
            sendMessage(session, createErrorMessage("SEND_PRIVATE_MESSAGE_ERROR", "发送私聊消息失败"));
        }
    }

    /**
     * 处理私聊消息已读
     */
    private void handleReadPrivateMessage(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        try {
            String sessionId = (String) wsRequest.getData();

            // 标记私聊消息为已读
            privateMessageService.markMessagesAsRead(sessionId, sessionInfo.getUserId());

            // 发送已读确认
            sendMessage(session, createSystemMessage("PRIVATE_MESSAGES_READ", "私聊消息已标记为已读"));

        } catch (Exception e) {
            log.error("处理私聊消息已读异常", e);
            sendMessage(session, createErrorMessage("READ_PRIVATE_MESSAGE_ERROR", "标记私聊已读失败"));
        }
    }

    /**
     * 处理心跳
     */
    private void handleHeartbeat(WebSocketSession session, UserSessionInfo sessionInfo) {
        sessionInfo.setLastHeartbeat(System.currentTimeMillis());
        sendMessage(session, createSystemMessage("HEARTBEAT_ACK", "心跳响应"));
    }

    /**
     * 向指定用户发送消息
     */
    public void sendToUser(String userId, Object message) {
        CopyOnWriteArraySet<WebSocketSession> sessions = USER_SESSIONS.get(userId);
        if (sessions != null) {
            sessions.forEach(session -> sendMessage(session, message));
        }
    }

    /**
     * 向群组广播消息
     */
    public void broadcastToGroup(String groupId, Object message, String excludeUserId) {
        CopyOnWriteArraySet<String> subscribers = GROUP_SUBSCRIBERS.get(groupId);
        if (subscribers != null) {
            subscribers.forEach(userId -> {
                if (!userId.equals(excludeUserId)) {
                    sendToUser(userId, message);
                }
            });
        }
    }

    /**
     * 发送消息到WebSocket会话
     */
    private void sendMessage(WebSocketSession session, Object message) {
        if (session.isOpen()) {
            try {
                String json = objectMapper.writeValueAsString(message);
                session.sendMessage(new TextMessage(json));
            } catch (IOException e) {
                log.error("发送WebSocket消息异常", e);
            }
        }
    }

    /**
     * 创建系统消息
     */
    private Map<String, Object> createSystemMessage(String type, String message) {
        return Map.of(
            "type", type,
            "message", message,
            "timestamp", System.currentTimeMillis()
        );
    }

    /**
     * 创建数据消息
     */
    private Map<String, Object> createDataMessage(String type, Object data) {
        return Map.of(
            "type", type,
            "data", data,
            "timestamp", System.currentTimeMillis()
        );
    }

    /**
     * 创建错误消息
     */
    private Map<String, Object> createErrorMessage(String errorCode, String errorMessage) {
        return Map.of(
            "type", "ERROR",
            "errorCode", errorCode,
            "errorMessage", errorMessage,
            "timestamp", System.currentTimeMillis()
        );
    }

    /**
     * WebSocket消息类
     */
    public static class WebSocketRequest {
        private String type;
        private Object data;

        // getters and setters
        public String getType() { return type; }
        public void setType(String type) { this.type = type; }
        public Object getData() { return data; }
        public void setData(Object data) { this.data = data; }
    }

    /**
     * 用户会话信息
     */
    public static class UserSessionInfo {
        private String userId;
        private String groupId;
        private String clientIp;
        private long connectTime;
        private long lastHeartbeat;

        // getters and setters
        public String getUserId() { return userId; }
        public void setUserId(String userId) { this.userId = userId; }
        public String getGroupId() { return groupId; }
        public void setGroupId(String groupId) { this.groupId = groupId; }
        public String getClientIp() { return clientIp; }
        public void setClientIp(String clientIp) { this.clientIp = clientIp; }
        public long getConnectTime() { return connectTime; }
        public void setConnectTime(long connectTime) { this.connectTime = connectTime; }
        public long getLastHeartbeat() { return lastHeartbeat; }
        public void setLastHeartbeat(long lastHeartbeat) { this.lastHeartbeat = lastHeartbeat; }
    }
}
